use std::alloc::AllocError;
use std::collections::HashMap;
use std::fmt::{Debug, Display, Formatter};
use std::hash::{BuildHasher, BuildHasherDefault};
use std::sync::atomic::{AtomicBool, Ordering};

use ahash::AHasher;
use bytestring::ByteString;
use once_cell::sync::{Lazy, OnceCell};
use parking_lot::{Mutex, RwLock};

use memory::Segment;

const SHARD_COUNT: usize = 64;

pub type BuildAHasher = BuildHasherDefault<AHasher>;

/// The Global SymbolMgr
static mut SYMBOL_MGR: Lazy<OnceCell<SymbolMgr<BuildAHasher>>> = Lazy::new(|| {
    OnceCell::with_value(SymbolMgr::create().expect("OutOfMemory when creating SymbolMgr"))
});

/// Thread safe
pub struct SymbolMgr<S> {
    shards: Vec<RwLock<Shard>>,
    hasher_builder: S,
}

struct Shard {
    shard_id: u16,
    data: Vec<Segment>,
    cur: *mut Segment,
    mapping: HashMap<String, RawSymbol>,
}

#[derive(Copy, Clone)]
union RawSymbol {
    info: RawSymbolInfo,
    // logic id
    value: u64,
}

/// Encode of raw symbol to u64
#[derive(Copy, Clone)]
struct RawSymbolInfo {
    shard_idx: u16,
    seg_idx: u16,
    seg_offset: u16,
    seg_length: u16,
}

#[derive(Copy, Clone, Eq, PartialEq, Hash, Ord, PartialOrd)]
pub struct Symbol(u64);

impl<S: BuildHasher + Default> SymbolMgr<S> {
    pub fn create() -> Result<Self, AllocError> {
        Self::create_with_hasher(S::default())
    }

    pub fn create_with_hasher(s: S) -> Result<Self, AllocError> {
        let mut shards = Vec::with_capacity(SHARD_COUNT);
        for i in 0..SHARD_COUNT {
            shards.push(RwLock::new(Shard::create(i as _)?));
        }
        Ok(Self {
            shards,
            hasher_builder: s,
        })
    }

    pub fn get(&mut self, key: &str) -> Result<Symbol, AllocError> {
        let slot_idx = self.hasher_builder.hash_one(key) % SHARD_COUNT as u64;
        let shard = &mut self.shards[slot_idx as usize];
        shard.write().get(key)
    }

    pub fn contains(&self, key: &String) -> bool {
        let slot_idx = self.hasher_builder.hash_one(key) % SHARD_COUNT as u64;
        let shard = &self.shards[slot_idx as usize];
        shard.read().contains(key)
    }
}

impl<S> SymbolMgr<S> {
    pub unsafe fn get_string_ref(&self, symbol: Symbol) -> ByteString {
        let raw = RawSymbol::from(symbol);
        debug_assert!(raw.info.shard_idx < SHARD_COUNT as u16, "shard_idx: {}", raw.info.shard_idx);
        let shard = &self.shards[raw.info.shard_idx as usize];
        shard.read().get_string_ref(raw)
    }

    pub fn get_string(&self, symbol: Symbol) -> ByteString {
        let raw = RawSymbol::from(symbol);
        let shard = unsafe {
            debug_assert!(raw.info.shard_idx < SHARD_COUNT as u16, "shard_idx: {}", raw.info.shard_idx);
            &self.shards[raw.info.shard_idx as usize]
        };
        shard.read().get_string(raw)
    }

    pub unsafe fn clear(&mut self) -> Result<(), AllocError> {
        for shard in &mut self.shards {
            shard.write().clear()?;
        }
        Ok(())
    }

    pub fn used_bytes(&self) -> usize {
        self.shards.iter().map(|s| s.read().used_bytes()).sum()
    }

    pub fn alloc_bytes(&self) -> usize {
        self.shards.iter().map(|s| s.read().alloc_bytes()).sum()
    }

    pub fn symbol_count(&self) -> usize {
        self.shards.iter().map(|s| s.read().symbol_count()).sum()
    }
}

impl Shard {
    pub fn create(shard_id: u16) -> Result<Self, AllocError> {
        let segment = unsafe { Segment::create()? };
        // segment in stack, put segment into vec and move to heap
        let data = vec![segment];
        // get the raw pointer on the heap
        let cur = &data[0] as *const Segment as *mut Segment;
        Ok(Self {
            shard_id,
            data,
            cur,
            mapping: Default::default(),
        })
    }

    fn get(&mut self, key: &str) -> Result<Symbol, AllocError> {
        if let Some(raw_symbol) = self.mapping.get(key) {
            unsafe {
                Ok(raw_symbol.value.into())
            }
        } else {
            unsafe {
                let raw_symbol = if (&*self.cur).has_left_space(key.len()) {
                    let idx = self.data.len() - 1;
                    let offset = (&*self.cur).offset();
                    (&mut *self.cur).write_uncheck(key.as_bytes());
                    RawSymbol::create(self.shard_id, idx, offset, key.len())
                } else {
                    let mut seg = Segment::create()?;
                    seg.write_uncheck(key.as_bytes());
                    self.data.push(seg);
                    self.cur = self.data.last().unwrap() as *const Segment as *mut Segment;
                    RawSymbol::create(self.shard_id, self.data.len() - 1, 0, key.len())
                };
                self.mapping.insert(key.into(), raw_symbol);
                Ok(raw_symbol.value.into())
            }
        }
    }

    /// # Unsafe
    /// ByteString may out of mem-store or at wrong pos if SymbolMgr after GC.
    #[inline]
    unsafe fn get_string_ref(&self, raw: RawSymbol) -> ByteString {
        unsafe {
            let seg = &self.data[raw.info.seg_idx as usize];
            let b = seg.slice_uncheck(raw.info.seg_offset as isize, raw.info.seg_length as usize);
            ByteString::from_bytes_unchecked(b)
        }
    }

    #[inline]
    fn get_string(&self, raw: RawSymbol) -> ByteString {
        unsafe {
            let seg = &self.data[raw.info.seg_idx as usize];
            let b = seg.slice_copy(raw.info.seg_offset as isize, raw.info.seg_length as usize);
            ByteString::from_bytes_unchecked(b)
        }
    }

    unsafe fn clear(&mut self) -> Result<(), AllocError> {
        self.mapping.clear();
        self.data.clear();
        let seg = Segment::create()?;
        self.data.push(seg);
        self.cur = &self.data[0] as *const Segment as *mut Segment;
        Ok(())
    }

    #[inline]
    fn contains(&self, key: &String) -> bool {
        self.mapping.contains_key(key)
    }

    fn used_bytes(&self) -> usize {
        self.data.iter().map(|s| s.offset()).sum::<isize>() as usize
    }

    fn alloc_bytes(&self) -> usize {
        self.data.iter().map(|s| s.length()).sum()
    }

    #[inline]
    fn symbol_count(&self) -> usize {
        self.mapping.len()
    }
}

#[inline]
pub fn to_symbol(key: &str) -> Result<Symbol, AllocError> {
    unsafe {
        let mgr = (&mut *SYMBOL_MGR).get_mut().expect("SymbolMgr not init");
        mgr.get(key)
    }
}


impl RawSymbol {
    #[inline]
    fn create(shard_idx: u16, seg_idx: usize, seg_offset: isize, seg_length: usize) -> Self {
        Self {
            info: RawSymbolInfo {
                shard_idx,
                seg_idx: seg_idx as _,
                seg_offset: seg_offset as _,
                seg_length: seg_length as _,
            }
        }
    }
}

impl From<Symbol> for RawSymbol {
    #[inline]
    fn from(v: Symbol) -> Self {
        RawSymbol { value: v.0 }
    }
}

impl From<u64> for Symbol {
    #[inline]
    fn from(v: u64) -> Self {
        Symbol(v)
    }
}

impl Symbol {

    #[inline]
    pub unsafe fn get_string_ref(self) -> ByteString {
        get_symbol_mgr().get_string_ref(self)
    }

    #[inline]
    pub fn get_string(self) -> ByteString {
        get_symbol_mgr().get_string(self)
    }

    #[inline]
    pub fn get_value(self) -> u64 {
        self.0
    }
}

impl Debug for Symbol {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        let rs = RawSymbol::from(*self);
        unsafe {
            write!(f, "{}_$_{}:{}:{}:{}", self.get_string_ref(), rs.info.shard_idx, rs.info.seg_idx, rs.info.seg_offset, rs.info.seg_length)
        }
    }
}

impl TryFrom<&String> for Symbol {
    type Error = AllocError;
    #[inline]
    fn try_from(v: &String) -> Result<Self, Self::Error> {
        to_symbol(v.as_str())
    }
}

impl TryFrom<ByteString> for Symbol {
    type Error = AllocError;
    #[inline]
    fn try_from(v: ByteString) -> Result<Self, Self::Error> {
        to_symbol(&*v)
    }
}

impl TryFrom<&ByteString> for Symbol {
    type Error = AllocError;
    #[inline]
    fn try_from(v: &ByteString) -> Result<Self, Self::Error> {
        to_symbol(&*v)
    }
}

impl TryFrom<String> for Symbol {
    type Error = AllocError;
    #[inline]
    fn try_from(v: String) -> Result<Self, Self::Error> {
        to_symbol(&v)
    }
}

impl TryFrom<&str> for Symbol {
    type Error = AllocError;
    #[inline]
    fn try_from(v: &str) -> Result<Self, Self::Error> {
        to_symbol(v)
    }
}

impl Display for Symbol {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        unsafe {
            write!(f, "{}", &*self.get_string_ref())
        }
    }
}

static SYMBOL_INIT_FLAG: AtomicBool = AtomicBool::new(false);
static SYMBOL_INIT_LOCK: Mutex<()> = Mutex::new(());

pub fn init_symbol_system() {
    let lock = SYMBOL_INIT_LOCK.lock();
    if SYMBOL_INIT_FLAG.compare_exchange(false, true, Ordering::SeqCst, Ordering::Acquire).is_ok() {
        unsafe {
            let _ = &*SYMBOL_MGR;
        }
    }
    drop(lock);
}

#[inline]
pub fn get_symbol_mgr() -> &'static SymbolMgr<BuildAHasher> {
    unsafe {
        &*SYMBOL_MGR.get().unwrap()
    }
}

#[inline]
pub fn get_symbol_mgr_mut() -> &'static mut SymbolMgr<BuildAHasher> {
    unsafe {
        &mut *SYMBOL_MGR.get_mut().unwrap()
    }
}


#[cfg(test)]
mod test {
    use std::collections::HashSet;
    use std::sync::Arc;
    use std::thread::spawn;

    use dashmap::DashSet;
    use parking_lot::Mutex;
    use rand::{Rng, thread_rng};

    use memory::DEFAULT_SEGMENT_SIZE;

    use crate::mgr::{BuildAHasher, get_symbol_mgr, get_symbol_mgr_mut, init_symbol_system, SHARD_COUNT, to_symbol};
    use crate::SymbolMgr;

    static SEQ_LOCK: Mutex<()> = Mutex::new(());

    fn lock_in_seq<F: FnOnce() -> ()>(f: F) {
        let lock = SEQ_LOCK.lock();
        f();
        unsafe {
            get_symbol_mgr_mut().clear().unwrap();
        }
        drop(lock);
    }

    #[test]
    fn test() {
        lock_in_seq(|| {
            init_symbol_system();
            let key = "abcd1235".to_string();
            for _ in 0..10 {
                let _symbol = to_symbol(&key).unwrap();
                // println!("symbol: {}", symbol);
            }
            let mgr = get_symbol_mgr();
            assert_eq!(1, mgr.symbol_count());
            assert_eq!(8, mgr.used_bytes());
            assert_eq!(DEFAULT_SEGMENT_SIZE * SHARD_COUNT, mgr.alloc_bytes());
        });
    }

    #[test]
    fn test1() {
        lock_in_seq(|| {
            init_symbol_system();
            do_rnd_test(10000, 100);
        });
    }

    #[test]
    fn test_many_segments() {
        lock_in_seq(|| {
            init_symbol_system();
            do_rnd_test(100_000, 10_000);
        });
    }

    #[test]
    fn test_multi_threads() {
        lock_in_seq(|| {
            init_symbol_system();
            do_test_multi_threads(10_000, 1_000, 10);
        });
    }

    #[test]
    fn test_multi_threads_1() {
        lock_in_seq(|| {
            init_symbol_system();
            do_test_multi_threads(100_000, 100_000, 32);
        });
    }

    fn do_rnd_test(loop_n: usize, max_n: usize) {
        let mut keys = HashSet::new();
        let mut rnd = thread_rng();
        for _ in 0..loop_n {
            let key = format!("long_long_long_key_{}", rnd.gen_range(0..max_n));
            let _symbol = to_symbol(&key).unwrap();
            keys.insert(key);
        }
        let mgr = get_symbol_mgr();
        print_diff(&keys, mgr);
        assert_eq!(keys.len(), mgr.symbol_count());
        let byte_used: usize = keys.iter().map(|s| s.len()).sum();
        assert_eq!(byte_used, mgr.used_bytes());
        println!("symbol_count: {}, bytes_used: {}, bytes_alloc: {}", keys.len(), byte_used, mgr.alloc_bytes());
    }

    fn print_diff(keys: &HashSet<String>, mgr: &SymbolMgr<BuildAHasher>) {
        mgr.shards.iter().for_each(|s| {
            s.read().mapping.keys().for_each(|k| {
                if !keys.contains(k) {
                    println!("key: {} in SymbolMgr, not in HashSet", k);
                }
            })
        });
        keys.iter().for_each(|k| {
            if !mgr.contains(k) {
                println!("key: {} in HashSet, not in SymbolMgr", k);
            }
        })
    }

    fn do_test_multi_threads(loop_n: usize, max_n: usize, thread_num: usize) {
        let keys = Arc::new(DashSet::new());
        let mut handles = vec![];
        for _ in 0..thread_num {
            let keys0 = keys.clone();
            let handle = spawn(move || {
                do_add_rnd_keys(keys0, loop_n, max_n);
            });
            handles.push(handle);
        }
        // wait all finish
        for handle in handles {
            handle.join().unwrap();
        }
        let mgr = get_symbol_mgr();
        // print_diff(&keys, mgr);
        assert_eq!(keys.len(), mgr.symbol_count());
        let byte_used: usize = keys.iter().map(|s| s.len()).sum();
        assert_eq!(byte_used, mgr.used_bytes());
        println!("symbol_count: {}, bytes_used: {}, bytes_alloc: {}", keys.len(), byte_used, mgr.alloc_bytes());
    }

    fn do_add_rnd_keys(keys: Arc<DashSet<String>>, loop_n: usize, max_n: usize) {
        let mut rnd = thread_rng();
        for _ in 0..loop_n {
            let key = format!("this_is_very_long_key_{}", rnd.gen_range(0..max_n));
            let _symbol = to_symbol(&key).unwrap();
            keys.insert(key);
        }
    }
}